#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
网络连接错误修复工具
帮助解决rembg模型下载失败的问题
"""

import os
import sys
import requests
from pathlib import Path


def check_internet_connection():
    """检查网络连接状态"""
    print("🔍 检查网络连接...")
    
    test_urls = [
        "https://www.baidu.com",
        "https://www.google.com",
        "https://github.com"
    ]
    
    for url in test_urls:
        try:
            response = requests.get(url, timeout=5)
            if response.status_code == 200:
                print(f"✅ 网络连接正常: {url}")
                return True
        except Exception as e:
            print(f"❌ 无法连接到 {url}: {e}")
    
    print("❌ 网络连接异常")
    return False


def configure_requests_settings():
    """配置requests库的网络设置"""
    print("🔧 配置网络请求设置...")
    
    # 设置更长的超时时间
    import socket
    socket.setdefaulttimeout(30)
    
    # 配置urllib3
    try:
        import urllib3
        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        print("✅ 禁用SSL警告")
    except ImportError:
        pass
    
    # 配置requests默认设置
    try:
        import requests.adapters
        from requests.packages.urllib3.util.retry import Retry
        
        # 配置重试策略
        retry_strategy = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504]
        )
        
        adapter = requests.adapters.HTTPAdapter(max_retries=retry_strategy)
        session = requests.Session()
        session.mount("http://", adapter)
        session.mount("https://", adapter)
        
        print("✅ 配置网络重试策略")
    except Exception as e:
        print(f"⚠️ 网络配置警告: {e}")


def download_models_manually():
    """手动下载模型文件"""
    print("📦 手动下载AI模型...")
    
    models_info = {
        'u2net': {
            'url': 'https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx',
            'filename': 'u2net.onnx',
            'size_mb': 176
        }
    }
    
    model_dir = Path.cwd() / "models"
    model_dir.mkdir(exist_ok=True)
    
    for model_name, info in models_info.items():
        model_file = model_dir / info['filename']
        
        if model_file.exists():
            print(f"✅ 模型已存在: {info['filename']}")
            continue
        
        print(f"📥 下载模型: {info['filename']} ({info['size_mb']}MB)")
        
        try:
            # 使用分片下载和重试机制
            response = requests.get(
                info['url'], 
                stream=True, 
                timeout=60,
                headers={'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'}
            )
            response.raise_for_status()
            
            total_size = int(response.headers.get('content-length', 0))
            downloaded = 0
            
            with open(model_file, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        
                        if total_size > 0:
                            progress = (downloaded / total_size) * 100
                            print(f"\r   进度: {progress:.1f}% ({downloaded // (1024*1024)}MB/{total_size // (1024*1024)}MB)", end='')
            
            print(f"\n✅ 下载完成: {info['filename']}")
            
        except Exception as e:
            print(f"\n❌ 下载失败: {e}")
            if model_file.exists():
                model_file.unlink()  # 删除不完整的文件


def fix_connection_error():
    """修复连接错误的主函数"""
    print("🔧 网络连接错误修复工具")
    print("=" * 50)
    
    # 1. 检查网络连接
    has_internet = check_internet_connection()
    print()
    
    # 2. 配置网络设置
    configure_requests_settings()
    print()
    
    # 3. 如果有网络，尝试下载模型
    if has_internet:
        try:
            download_models_manually()
        except Exception as e:
            print(f"❌ 模型下载失败: {e}")
    else:
        print("⚠️ 无网络连接，跳过模型下载")
    print()
    
    print("🎉 修复完成！")
    print()
    print("📝 解决方案说明:")
    print("1. ✅ 已配置网络重试策略")
    print("2. ✅ 已尝试下载AI模型")
    print()
    print("🚀 现在可以重新运行您的程序了")